package com.yihaodian.index.solr.ext;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.log4j.Logger;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.DocsEnum;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexReaderContext;
import org.apache.lucene.index.MultiFields;
import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.docvalues.FloatDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.BytesRef;
import org.apache.solr.schema.SchemaField;

public class MemFloatSource extends ValueSource {

	private static Logger log = Logger.getLogger(MemFloatSource.class);
	private MemDataProvider provider;

	private String indexName;
	private SchemaField field;

	private final SchemaField keyField;
	private final float defVal;
	private final String dataDir;

	private int maxDocs;
	private float[] values;

	// key -> docs
	private Map<String, List<Integer>> keyDocs;

	public MemFloatSource(MemDataProvider provider, String indexName,
			SchemaField field, SchemaField keyField, float defVal,
			String datadir) {
		this.indexName = indexName;
		this.field = field;
		this.keyField = keyField;
		this.defVal = defVal;
		this.dataDir = datadir;
		this.values = null;

		this.provider = provider;
	}

	@Override
	public String description() {
		return "float(" + field + ')';
	}

	@Override
	public boolean equals(Object o) {
		if (o.getClass() != MemFloatSource.class)
			return false;
		MemFloatSource other = (MemFloatSource) o;
		return this.indexName.equals(other.indexName)
				&& this.field.getName().equals(other.field.getName())
				&& this.keyField.getName().equals(other.keyField.getName())
				&& this.defVal == other.defVal
				&& this.dataDir.equals(other.dataDir);
	}

	@Override
	public FunctionValues getValues(Map context,
			AtomicReaderContext readerContext) throws IOException {
		final int off = readerContext.docBase;
		IndexReaderContext topLevelContext = ReaderUtil
				.getTopLevelContext(readerContext);

		final float[] arr = getCachedFloats(topLevelContext.reader());
		return new FloatDocValues(this) {
			@Override
			public float floatVal(int doc) {
				if(arr == null || doc + off > arr.length - 1) {
					return defVal;
				}
				return arr[doc + off];
			}

			@Override
			public Object objectVal(int doc) {
				return floatVal(doc); // TODO: keep track of missing values
			}
		};
	}

	@Override
	public int hashCode() {
		return MemFloatSource.class.hashCode() + field.getName().hashCode();
	}

	private final float[] getCachedFloats(IndexReader reader) {

		if (this.values == null) {
			loadValues(reader, true);
		}

		return values;
	}

	public synchronized void loadValues(IndexReader reader, boolean force) {
		if (this.keyDocs == null || force) {
			loadKeyDocs(reader, true);
		}
		if (keyDocs == null) {
			return;
		}

		if (this.values == null || this.values.length != maxDocs || force) {			
			loadValues();
		}
	}

	/**
	 * Load all key from docs
	 * @param reader
	 * @param force
	 */
	private synchronized void loadKeyDocs(IndexReader reader, boolean force) {
		if (!force && this.keyDocs != null) {
			return;
		}
		// load key -> docs mapping from index
		try {
			Map<String, List<Integer>> newData = new HashMap<String, List<Integer>>(); 
			
			Terms t =  MultiFields.getTerms(reader, keyField.getName());
			if(t == null) {
				log.error("Failed to get terms, reader: " + reader + ", keyField: " + keyField);
				return;
			}
			TermsEnum te = t.iterator(null);

			DocsEnum de = null;
			BytesRef term;
			while (true) {
				term = te.next();
				if (term == null) {
					break;
				}
				String key = term.utf8ToString();

				de = te.docs(null, de);
				List<Integer> docs = new ArrayList<Integer>();
				int doc = 0;
				while ((doc = de.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
					docs.add(doc);
				}
				if (!docs.isEmpty()) {
					newData.put(key, docs);
				}
			}
			
			this.keyDocs = newData;
			this.maxDocs = reader.maxDoc();
		} catch (IOException e) {
			log.error("Failed to load key-docs mapping from index.", e);			
		}

	}

	
	public synchronized void loadValues() {
		
		if(keyDocs == null) {
			return;
		}
		
		float[] newValues = new float[maxDocs];
		if (defVal != 0.0) {
			Arrays.fill(newValues, defVal);
		}

		for (Entry<String, List<Integer>> entry : keyDocs.entrySet()) {

			Float value = provider.get(entry.getKey());
			if(value == null || value.equals(defVal)) {
				continue;
			}

			List<Integer> docs = entry.getValue();
			if (docs != null) {
				for (int doc : docs) {
					newValues[doc] = value;
				}
			}

		}

		this.values = newValues;
	}
	
	public void refresh(String key, float value) {
		if(this.values == null) {
			return;
		}
		
		List<Integer> docs = keyDocs.get(key);
		if(docs == null) {
			return;
		}
		for (Integer doc : docs) {
			if(doc > this.values.length - 1) {
				log.error("docs is not refreshed yet, doc: " + doc + ", key: " + key + ", value: " + value);
				continue;
			}
			this.values[doc] = value;
		}
	}


	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		
		sb.append("MemFloatSource, keyDocs: ").append(keyDocs);
		sb.append(", values: ");
		if(values != null) {
			for(int i = 0; i < values.length; ++ i) {
				sb.append(values[i]).append(", ");
			}
		} else {
			sb.append("null");
		}
		return sb.toString();
	}
}
